from ultralytics import YOLO
import os
import cv2


def calibrate_model():
    # 加载最佳模型
    model = YOLO("runs/detect/train-yolo8l-150t-alldata/weights/best.pt")

    # 标定数据文件夹
    calibration_folder = os.path.join("datasets", "calibration")

    # 检查标定数据文件夹是否存在
    if not os.path.exists(calibration_folder):
        print("未找到标定数据文件夹！")
        return

    # 获取标定数据文件列表
    image_extensions = [".png", ".jpg", ".jpeg", ".bmp", ".tiff"]
    calibration_data = []

    for root, _, files in os.walk(calibration_folder):
        for file in files:
            # 只处理图片文件
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_path = os.path.join(root, file)
                # 获取对应的标签文件名（替换扩展名）
                label_file = os.path.splitext(file)[0] + ".txt"
                label_path = os.path.join(root, label_file)

                # 验证标签文件是否存在
                if os.path.exists(label_path):
                    calibration_data.append((image_path, label_path))
                else:
                    print(f"警告：图片 {image_path} 没有对应的标签文件 {label_path}")

    print(f"找到 {len(calibration_data)} 对有效标定数据")

    # 创建临时数据配置文件
    if calibration_data:
        # 创建临时数据配置文件
        temp_yaml_path = "datasets/temp_calibration.yaml"

        # 获取标定文件夹的绝对路径
        abs_calibration_folder = os.path.abspath(calibration_folder)

        with open(temp_yaml_path, "w") as f:
            # 使用绝对路径避免相对路径问题
            f.write(f"train: {abs_calibration_folder}\n")
            f.write(f"val: {abs_calibration_folder}\n")

            # 从原始配置文件中复制类别信息
            try:
                with open("datasets/data.yaml", "r") as original_f:
                    for line in original_f:
                        if line.startswith("nc") or line.startswith("names"):
                            f.write(line)
            except Exception as e:
                print(f"警告：无法复制原始配置文件中的类别信息: {e}")
                f.write("nc: 1\n")  # 默认设置为1类
                f.write("names: ['object']\n")

        # 训练参数
        epochs = 10  # 微调的轮数，可以根据需要调整
        imgsz = 640
        device = "cpu"  # 若无GPU用 device="cpu"

        # 开始微调
        model.train(
            data=temp_yaml_path,  # 使用临时配置文件
            epochs=epochs,
            imgsz=imgsz,
            device=device,
            resume=False,  # 不恢复之前的训练状态
        )

        # 获取最新训练结果的保存目录（假设是最后一个文件夹）
        train_dirs = sorted(
            [d for d in os.listdir("runs/detect") if d.startswith("train")],
            key=lambda x: int(x[5:]) if x[5:].isdigit() else 0,
            reverse=True,
        )
        if not train_dirs:
            print("未找到训练结果目录！")
            return

        latest_train_dir = os.path.join("runs/detect", train_dirs[0])
        weights_dir = os.path.join(latest_train_dir, "weights")

        # 源文件路径（best.pt）
        source_path = os.path.join(weights_dir, "best.pt")

        # 目标文件路径（自定义名称）
        target_name = "calibrated_best.pt"  # 自定义名称
        target_path = os.path.join(weights_dir, target_name)

        # 重命名文件
        if os.path.exists(source_path):
            os.rename(source_path, target_path)
            print(f"模型已重命名并保存到 {target_path}")
        else:
            print(f"源文件不存在: {source_path}")

        # 清理临时文件
        if os.path.exists(temp_yaml_path):
            os.remove(temp_yaml_path)
    else:
        print("未找到有效标定数据！")


if __name__ == "__main__":
    calibrate_model()
